#include "svd.h"
#include "gnu_timer.h"
#include <stdio.h>
#include <math.h>
#include "marker.h"

#ifndef MAX_ITER
#define MAX_ITER 30
#endif

static struct colpair_t ColPairs;
static int NT = 1;
static REAL orth_chk;

REAL inner_product(REAL *a, REAL *b, int len)
{
    int i;
    REAL sum = 0;
    for (i = 0; i < len; ++i)
    {
        sum += a[i] * b[i];
    }
    return sum;
}

void update_norm(struct matrix_t mat)
{
    int i;
#pragma omp parallel for num_threads(NT)
    for (i = 0; i < mat.w; ++i)
    {
        mat.col_idx[i] = i;
    }
#pragma omp parallel for num_threads(NT)
    for (i = 0; i < mat.w; ++i)
    {
        mat.norms[i] = inner_product(&mat.in[i * mat.h], &mat.in[i * mat.h], mat.h);
    }
}

inline static void
sort(REAL *data, int *idx, int len)
{
    int cnt, i;
    for (cnt = 0; cnt < len; ++cnt)
#pragma omp parallel for num_threads(NT)
        for (i = cnt % 2; i < len; i += 2)
        {
            if (data[i] < data[i + 1])
            {
                swap_data(&data[i], &data[i + 1]);
                swap_index(&idx[i], &idx[i + 1]);
            }
        }
}

void sort_norm(struct matrix_t mat)
{
    int num_cols = mat.w;
    int *sorted_idx = mat.col_idx;
    REAL *norms = mat.norms;
    sort(norms, sorted_idx, num_cols);
}

static inline bool
rotate_pair(REAL *va[2], REAL *vb[2], REAL *norms[2], int vlen)
{
    int i;
    REAL product = inner_product(va[0], va[1], vlen);
    // printf("%f\n", product);
    if (fabs(product) < THRESHOLD)
    {
        return true;
    }

    REAL tau = (*norms[0] - *norms[1]) / (2 * product);
    REAL tan = sign(tau) / (fabs(tau) + sqrt(1 + tau * tau));
    REAL cos = 1 / sqrt(1 + tan * tan);
    REAL sin = cos * tan;

    orth_chk = orth_chk < product ? product : orth_chk;
    *norms[0] = *norms[0] + product * tan;
    *norms[1] = *norms[1] - product * tan;

    for (i = 0; i < vlen; i++)
    {
        REAL ele0, ele1;
        ele0 = va[0][i] * cos + va[1][i] * sin;
        ele1 = va[1][i] * cos - va[0][i] * sin;
        va[0][i] = ele0;
        va[1][i] = ele1;
        ele0 = vb[0][i] * cos + vb[1][i] * sin;
        ele1 = vb[1][i] * cos - vb[0][i] * sin;
        vb[0][i] = ele0;
        vb[1][i] = ele1;
    }
    return false;
}



void switch_pair_trivial(int *aidx, int* bidx, int* aout, int* bout, int len, int phase)
{
    int i, swap_idx, tmp;
    swap_idx = len - 1 - phase / 2;
    // index a mostly remains the same
    memcpy(aout, aidx, sizeof(int)*len);
    tmp = bidx[0];
#pragma omp parallel for num_threads(NT)
    for (i = 0; i < len - 1; i++)
    {
        bout[i] = bidx[i + 1];
    }
    bout[len-1] = tmp;
    swap_index(&aout[swap_idx], &bout[(swap_idx+len-1)%len]);
}

void reset_pair_trivial()
{
    int i, num_cols = ColPairs.len;

#pragma omp parallel for num_threads(NT)
    for (i = 0; i < num_cols * 2; ++i)
    {
        if (i % 2 == 0)
            ColPairs.colidx_a[i / 2] = i;
        else
            ColPairs.colidx_b[i / 2] = i;
    }
}



#ifdef UNROLL

// 实现分层ring序列，将列分组，按组配对，组间ring序列，组对内使用最佳串行序列
// 这是一个放大粒度的操作，组内正好能放到一个L1cache里，所以用串行最快序列，由一个线程完成计算即可

void init_pair(int col_len)
{
    ColPairs.len = col_len;
    ColPairs.iters = (col_len + NT - 1) / NT;
    ColPairs.colidx_a = (int *)malloc(sizeof(int) * ColPairs.iters * col_len * 2);
    ColPairs.colidx_b = (int *)malloc(sizeof(int) * ColPairs.iters * col_len * 2);
}

// 值得注意的是，这个分块方法似乎不能采用最均匀的负载均衡方法，
static inline void
query_pair(int iter, int wid, int block_idx, int* aidx, int* bidx)
{
    const int limit = ColPairs.len % NT == 0? ColPairs.len : ColPairs.len % NT;
    if (block_idx < limit)
        wid += ColPairs.iters * block_idx;
    else
        wid += (ColPairs.iters-1) * block_idx + limit;
    *aidx = ColPairs.colidx_a[iter*ColPairs.len+wid];
    *bidx = ColPairs.colidx_b[iter*ColPairs.len+wid];
}

void reset_pair()
{
    reset_pair_trivial();
    for (int i=0; i<ColPairs.iters * 2 - 1; ++i)
    {
        int* next_aidx = &ColPairs.colidx_a[ColPairs.len*(i+1)];
        int* next_bidx = &ColPairs.colidx_b[ColPairs.len*(i+1)];
        int* this_aidx = &ColPairs.colidx_a[ColPairs.len*i];
        int* this_bidx = &ColPairs.colidx_b[ColPairs.len*i];
        switch_pair_trivial(this_aidx, this_bidx, next_aidx, next_bidx, ColPairs.len, i);
    }
}

// 我们希望这里的phase就是外面的base
void switch_pair(int phase, int step)
{
    const int len = ColPairs.len;
    const int max_phase = ColPairs.iters*2;
    int *from_aidx, *from_bidx, *to_aidx, *to_bidx;
    // 首先将后几phase前移step个phase
    for (int i=step; i<max_phase; i++)
    {
        from_aidx = &ColPairs.colidx_a[len*(i)];
        from_bidx = &ColPairs.colidx_b[len*(i)];
        to_aidx   = &ColPairs.colidx_a[len*(i-step)];
        to_bidx   = &ColPairs.colidx_b[len*(i-step)];
        memcpy(to_aidx, from_aidx, sizeof(int)*len);
        memcpy(to_bidx, from_bidx, sizeof(int)*len);
    }
    // 在生成那几个step的新phase
    for (int i=0; i<step; ++i)
    {
        from_aidx = &ColPairs.colidx_a[len*(max_phase - step + i - 1)];
        from_bidx = &ColPairs.colidx_b[len*(max_phase - step + i - 1)];
        to_aidx   = &ColPairs.colidx_a[len*(max_phase - step + i)];
        to_bidx   = &ColPairs.colidx_b[len*(max_phase - step + i)];
        switch_pair_trivial(from_aidx, from_bidx, to_aidx, to_bidx, len, max_phase-1+phase+i);
    }
}

// 这部分接受3列，分别对两列进行旋转然后与另一列进行内积计算工作
// 其中va[0,1]:需要旋转的两列， va[0,2]:需要算内积的两列
static inline REAL 
rotate_blocked_pair_stable(REAL *va[3], REAL *vb[2], REAL *norms[2], REAL product, int vlen)
{
    int i;
    REAL next_product=.0;
    if (fabs(product) < THRESHOLD)
    {
        return inner_product(va[1], va[2], vlen);
    }

    REAL tau = (*norms[0] - *norms[1]) / (2 * product);
    REAL tan = sign(tau) / (fabs(tau) + sqrt(1 + tau * tau));
    REAL cos = 1 / sqrt(1 + tan * tan);
    REAL sin = cos * tan;

    orth_chk = orth_chk < product ? product : orth_chk;
    *norms[0] = *norms[0] + product * tan;
    *norms[1] = *norms[1] - product * tan;

    for (i = 0; i < vlen; i++)
    {
        REAL ele0, ele1;
        ele0 = va[0][i] * cos + va[1][i] * sin;
        ele1 = va[1][i] * cos - va[0][i] * sin;
        va[0][i] = ele0;
        va[1][i] = ele1;
        next_product += va[2][i]*ele1;
        ele0 = vb[0][i] * cos + vb[1][i] * sin;
        ele1 = vb[1][i] * cos - vb[0][i] * sin;
        vb[0][i] = ele0;
        vb[1][i] = ele1;
    }
    return next_product;
}

// 这部分仅完成给定workload的旋转
static inline void 
rotate_blocked_pair_exit(REAL *va[2], REAL *vb[2], REAL *norms[2], REAL product, int vlen)
{
    int i;
    if (fabs(product) < THRESHOLD)
    {
        return; 
    }

    REAL tau = (*norms[0] - *norms[1]) / (2 * product);
    REAL tan = sign(tau) / (fabs(tau) + sqrt(1 + tau * tau));
    REAL cos = 1 / sqrt(1 + tan * tan);
    REAL sin = cos * tan;

    orth_chk = orth_chk < product ? product : orth_chk;
    *norms[0] = *norms[0] + product * tan;
    *norms[1] = *norms[1] - product * tan;

    for (i = 0; i < vlen; i++)
    {
        REAL ele0, ele1;
        ele0 = va[0][i] * cos + va[1][i] * sin;
        ele1 = va[1][i] * cos - va[0][i] * sin;
        va[0][i] = ele0;
        va[1][i] = ele1;
        ele0 = vb[0][i] * cos + vb[1][i] * sin;
        ele1 = vb[1][i] * cos - vb[0][i] * sin;
        vb[0][i] = ele0;
        vb[1][i] = ele1;
    }
    return;
}

#define Column(mat, idx) &mat.in[(idx)*mat.h]
// 计算第一个上三角workload
// 计算完后 step 1
static bool 
rotate_block_warmup(struct matrix_t A, struct matrix_t V, int blk_idx, int blk_size)
{
    int i,j;
    int phase, workload;
    int aidx, bidx, cidx;
    int arank, brank, crank, drank;
    REAL *Acol[3];
    REAL *Vcol[2];
    REAL *norms[2];
    REAL last_product;
    bool is_orth = true;
    for (i=0; i<blk_size; i++)
    {
        last_product = 0.0;
        for (j=i; j>=0; j--)
        {
            phase = i - j;
            workload = j;
            query_pair(phase, workload, blk_idx, &arank, &brank);
            aidx = A.col_idx[arank];
            bidx = A.col_idx[brank];
            Acol[0] = Column(A, aidx);
            Acol[1] = Column(A, bidx);
            Vcol[0] = Column(V, aidx);
            Vcol[1] = Column(V, bidx);
            norms[0] = &A.norms[arank];
            norms[1] = &A.norms[brank];
            if ( phase == 0 && workload == 0)
            {
                is_orth &= rotate_pair(Acol, Vcol, norms, A.h);
                // printf("[%d,%d] - %d %d\n", phase, workload, arank, brank);
                continue;
            }
            else if( phase == 0 )
            {
                last_product = inner_product(Acol[0], Acol[1], A.h);
                // printf("[%d,%d] = %f\n", arank, brank, last_product);
                is_orth &= last_product < THRESHOLD;
            }
            
            if( workload == 0 )
            {
                rotate_blocked_pair_exit(Acol, Vcol, norms, last_product, A.h);
                // printf("[%d,%d] - %d %d\n", phase, workload, arank, brank);
            }
            else
            {
                query_pair(phase+1, workload-1, blk_idx, &crank, &drank);
                if(crank == arank || drank == arank) //short-curcuit in theory
                {
                    SWAP(REAL*, Acol[0], Acol[1]);
                    SWAP(REAL*, Vcol[0], Vcol[1]);
                    SWAP(REAL*, norms[0], norms[1]);

                    SWAP(int, brank, arank);
                }
                if(crank == brank || crank == arank) crank = drank;
                // printf("[%d,%d] - %d %d %d\n", phase, workload, arank, brank, crank);
                cidx = A.col_idx[crank];
                Acol[2] = Column(A, cidx);
                last_product = rotate_blocked_pair_stable(Acol, Vcol, norms, last_product, A.h);
                // printf("[%d,%d] = %f\n", brank, crank, last_product);
                is_orth &= last_product < THRESHOLD;
            }
        }
    }
    return is_orth;
}

static bool 
rotate_block_stable(struct matrix_t A, struct matrix_t V, int blk_idx, int blk_size, int step, int base)
{
    int i,j;
    int phase, workload;
    int aidx, bidx, cidx;
    int arank, brank, crank, drank;
    REAL *Acol[3];
    REAL *Vcol[2];
    REAL *norms[2];
    REAL last_product;
    bool is_orth = true;
    for (i=0; i < step; ++i)
    {
        last_product = 0.0;
        for (j=blk_size-1; j>=0; j--)
        {
            phase = blk_size-1 - j + i;
            workload = j;
            if (phase + base > (ColPairs.len-1)*2) break;
            query_pair(phase, workload, blk_idx, &arank, &brank);
            aidx = A.col_idx[arank];
            bidx = A.col_idx[brank];
            Acol[0] = Column(A, aidx);
            Acol[1] = Column(A, bidx);
            Vcol[0] = Column(V, aidx);
            Vcol[1] = Column(V, bidx);
            norms[0] = &A.norms[arank];
            norms[1] = &A.norms[brank];
            if( workload == blk_size-1 )
            {
                last_product = inner_product(Acol[0], Acol[1], A.h);
                // printf("[%d,%d]=%f\n",arank, brank,last_product);
                is_orth &= last_product < THRESHOLD;
            }
            
            if( workload == 0 || phase + base == (ColPairs.len-1)*2)
            {
                // printf("[%d,%d] - %d %d\n", phase, workload, arank, brank);
                rotate_blocked_pair_exit(Acol, Vcol, norms, last_product, A.h);
            }
            else
            {
                query_pair(phase + 1, workload - 1, blk_idx, &crank, &drank);
                if(crank == arank || drank == arank) //short-curcuit in theory
                {
                    SWAP(REAL*, Acol[0], Acol[1]);
                    SWAP(REAL*, Vcol[0], Vcol[1]);
                    SWAP(REAL*, norms[0], norms[1]);

                    SWAP(int, brank, arank);
                }
                if(crank == arank || crank == brank) crank = drank;
                // printf("[%d,%d] - %d %d %d\n", phase, workload, arank, brank, crank);
                cidx = A.col_idx[crank];
                Acol[2] = Column(A, cidx);
                last_product = rotate_blocked_pair_stable(Acol, Vcol, norms, last_product, A.h);
                // printf("[%d,%d]=%f\n",brank, crank,last_product);
                is_orth &= last_product < THRESHOLD;
            }
        }
    }
    return is_orth;
}

#undef Column

bool grouped_rotate(struct matrix_t A, struct matrix_t V)
{
    int i, phase=0;
    const int limit = ColPairs.len % NT == 0? ColPairs.len : ColPairs.len % NT;
    int step = (limit < ColPairs.len) ? ColPairs.iters - 1 : ColPairs.iters;
    bool all_orth = true;
        //    for (int i=0; i<ColPairs.iters*2; i++)
        // {
        //     for(int j=0; j<ColPairs.len; j++)
        //         printf("%d ", ColPairs.colidx_a[ColPairs.len*i+j]);
        //     printf("\n");
        //     for(int j=0; j<ColPairs.len; j++)
        //         printf("%d ", ColPairs.colidx_b[ColPairs.len*i+j]);
        //     printf("\n\n");
        // }
        // printf("\n\n\n"); 
    #pragma omp parallel for num_threads(NT) reduction(&:all_orth)
    for (i=0; i<NT; i++)
    {
        int blk_idx = i;
        int blk_size = (blk_idx < limit) ? ColPairs.iters : ColPairs.iters - 1;
        all_orth &= rotate_block_warmup(A, V, blk_idx, blk_size);
    }
    switch_pair(phase, 1);
    phase += 1;
        
        //     for (int i=0; i<ColPairs.iters*2; i++)
        // {
        //     for(int j=0; j<ColPairs.len; j++)
        //         printf("%d ", ColPairs.colidx_a[ColPairs.len*i+j]);
        //     printf("\n");
        //     for(int j=0; j<ColPairs.len; j++)
        //         printf("%d ", ColPairs.colidx_b[ColPairs.len*i+j]);
        //     printf("\n\n");
        // }
        // printf("\n\n\n");
        
    while(phase < A.w - 1)
    {           
        // printf("step=%d\n", ColPairs.iters);
        #pragma omp parallel for num_threads(NT) reduction(&:all_orth)
        for (i=0; i<NT; i++)
        {
            int blk_idx = i;
            int blk_size = (blk_idx < limit) ? ColPairs.iters : ColPairs.iters - 1;
            all_orth &= rotate_block_stable(A, V, blk_idx, blk_size, step, phase);
        }
        switch_pair(phase, step);
        //         for (int i=0; i<ColPairs.iters*2; i++)
        // {
        //     for(int j=0; j<ColPairs.len; j++)
        //         printf("%d ", ColPairs.colidx_a[ColPairs.len*i+j]);
        //     printf("\n");
        //     for(int j=0; j<ColPairs.len; j++)
        //         printf("%d ", ColPairs.colidx_b[ColPairs.len*i+j]);
        //     printf("\n\n");
        // }
        // printf("\n\n\n");
        phase += step; 
        // printf("phase=%d\n", phase);
    }
    return all_orth;
}

#else

void init_pair(int col_len)
{
    ColPairs.len = col_len;
    ColPairs.iters = 1;
    ColPairs.colidx_a = (int *)malloc(sizeof(int) * col_len);
    ColPairs.colidx_b = (int *)malloc(sizeof(int) * col_len);
}


void reset_pair()
{
    reset_pair_trivial();
}

void switch_pair(int phase)
{
    switch_pair_trivial(ColPairs.colidx_a, ColPairs.colidx_b, ColPairs.colidx_a, ColPairs.colidx_b, ColPairs.len, phase);
}

bool grouped_rotate(struct matrix_t A, struct matrix_t V)
{
    int phase, i;
    bool all_orth = true;
    for (phase = 0; phase < A.w - 1; phase++)
    {
#pragma omp parallel for num_threads(NT) reduction(&:all_orth)
        for (i = 0; i < ColPairs.len; i++)
        {
            int pair_a = ColPairs.colidx_a[i];
            int pair_b = ColPairs.colidx_b[i];
            int idx_a = A.col_idx[pair_a];
            int idx_b = A.col_idx[pair_b];
            REAL *cols_in_A[2], *cols_in_V[2], *norms_ptr[2];

            // printf("[%d,%d] = ", pair_a, pair_b);
            cols_in_A[0] = &A.in[idx_a * A.h];
            cols_in_A[1] = &A.in[idx_b * A.h];
            cols_in_V[0] = &V.in[idx_a * V.h];
            cols_in_V[1] = &V.in[idx_b * V.h];
            norms_ptr[0] = &A.norms[pair_a];
            norms_ptr[1] = &A.norms[pair_b];

            all_orth &= rotate_pair(cols_in_A, cols_in_V, norms_ptr, A.h);
        }
        //printf("phase %d fin\n", phase);
        // SWAP(REAL *, A->in, A->out);
        switch_pair(phase);
    }
    return all_orth;
}

#endif

int format_output(struct matrix_t A, struct matrix_t V, REAL *Sigma)
{
    int i, j, nonz = 0;
    int *sorted_idx = (int *)malloc(sizeof(int) * A.w);

#pragma omp parallel for num_threads(NT)
    for (i = 0; i < A.w; i++)
        sorted_idx[i] = i;
    sort(A.norms, sorted_idx, A.w);

#pragma omp parallel for num_threads(NT)
    for (i = 0; i < A.w; i++)
        for (j = 0; j < A.h; j++)
            A.out[i * A.h + j] = A.in[sorted_idx[i] * A.h + j];
#pragma omp parallel for num_threads(NT)
    for (i = 0; i < V.w; i++)
        for (j = 0; j < V.h; j++)
            V.out[i * V.h + j] = V.in[sorted_idx[i] * V.h + j];
#pragma omp parallel for reduction(+:nonz) num_threads(NT)
    for (i = 0; i < A.w; i++)
    {
        Sigma[i] = sqrt(A.norms[i]);
        if (Sigma[i] > THRESHOLD)
            nonz+=1;
    }
    return nonz;
}

// 我们期望U只有out数组
void normalize(struct matrix_t A, struct matrix_t U, REAL *Sigma, int nonz)
{
    int i, j;
    int h = A.h;
    for (i = 0; i < nonz; i++)
    {
        if (Sigma[i] > THRESHOLD)
        {
#pragma omp parallel for num_threads(NT)
            for (j = 0; j < h; j++)
                U.out[i * h + j] = A.out[i * h + j] / sqrt(Sigma[i]);
        }
        else
        {
            REAL *Acol = &A.out[i * h + j];
#pragma omp parallel for num_threads(NT)
            for (j = 0; j < h; j++)
            {
                U.out[i * h + j] = A.out[i * h + j];
               //printf("Thread %d reporting\n", j);
            }
        }
    }
}

int main(int argc, char **argv)
{
    int i, num_iter, nonz;
    double start, end, start1, end1;
    struct matrix_t A, V, U;
    char *in_file;
    int *idx_buffer;
    REAL *Sigma;

    assert(argc == 3 && "usage：svd <num_threads> <*.mtx file>");
    NT = atoi(argv[1]);
    in_file = argv[2];
    matrix_from_file(&A, in_file);
    init_matrix(&V, A.w, A.w, ROTATE);
    init_matrix(&U, A.h, A.w, OUTPUT);

    idx_buffer = (int *)malloc(sizeof(int) * A.w);
    Sigma = (REAL *)malloc(sizeof(REAL) * A.w);
    memset(Sigma, 0, sizeof(REAL) * A.w);

    assert(A.w % 2 == 0 && "assumes num_cols = odd in order to simplify the problem");
    init_pair(A.w/2);
    //初始化旋转矩阵为单位1
    for (i = 0; i < V.w; i++)
        V.in[i * V.w + i] = 1.0;

LIKWID_MARKER_INIT;
LIKWID_MARKER_THREADINIT;
    num_iter = 0;
    start = mwtime();
LIKWID_MARKER_START("Compute");
CALC_BALANCER_START
    update_norm(A);
    while (num_iter < MAX_ITER)
    {
        orth_chk = 0;
        reset_pair();
        
        // for (int i=0; i<ColPairs.iters*2; i++)
        // {
        //     for(int j=0; j<ColPairs.len; j++)
        //         printf("%d ", ColPairs.colidx_a[ColPairs.len*i+j]);
        //     printf("\n");
        //     for(int j=0; j<ColPairs.len; j++)
        //         printf("%d ", ColPairs.colidx_b[ColPairs.len*i+j]);
        //     printf("\n\n");
        // }
        // printf("\n\n\n");
        
        sort_norm(A);
        start1 = mwtime();
        if (grouped_rotate(A, V))
            break;
        end1 = mwtime();
        printf(" iter %d residual %f in ms: %f\n", num_iter, orth_chk, end1 - start1);
//        getchar();
        num_iter++;
    }
CALC_BALANCER_END
LIKWID_MARKER_STOP("Compute");
    nonz = format_output(A, V, Sigma);
    normalize(A, U, Sigma, nonz);
    end = mwtime();
    matrix_to_file(U, "U.txt");
    matrix_to_file(V, "V.txt");
    vector_to_file(Sigma, A.w, "sigma.txt");
    printf("time elapses : %f\n", end - start);
    printf("iterations : %d\n", num_iter);
    printf("residual: %.10f\n", orth_chk);
LIKWID_MARKER_CLOSE;
}
